{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dataset_io import *\n",
    "from color_channels import equalize_hist\n",
    "import os\n",
    "import cv2\n",
    "import numpy as np\n",
    "import random as ra\n",
    "from tqdm import tqdm\n",
    "from profiler import *\n",
    "import random\n",
    "from PIL import Image, ImageCms\n",
    "from joblib import load"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def imread_resized(scene_name, filename, label_resolution):\n",
    "    filepath=get_filepath(scene_name, filename)\n",
    "    image = cv2.imread(filepath)\n",
    "    sh=image.shape\n",
    "    sh=(int(sh[1]/3), int(sh[0]/3))\n",
    "    return cv2.resize(image, sh)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def convert_to_gradient_channels(img):\n",
    "    channels=[]\n",
    "    \n",
    "    img = cv2.UMat(img)\n",
    "    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY )\n",
    "    channels.append(cv2.UMat(gray))\n",
    "    \n",
    "    gray = cv2.GaussianBlur(gray, (3,3), sigmaX=0, sigmaY=0)\n",
    "    channels.append(cv2.UMat(gray))\n",
    "    \n",
    "    gray = cv2.equalizeHist(gray) #experimental\n",
    "    gradX = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=3, scale=1 )\n",
    "    gradY = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=3, scale=1 )\n",
    "    mag, ori = cv2.cartToPolar(gradX, gradY, angleInDegrees=True)\n",
    "\n",
    "    ori_ranges = [(0,60), (60,120), (120,180), (180,240), (240,300), (300,360)]\n",
    "    oris=[cv2.inRange(ori,l,r) for (l,r) in ori_ranges]\n",
    "    mags= [cv2.copyTo(mag, ori_x) for ori_x in oris]\n",
    "    channels.extend([gray,mag]+mags)\n",
    "    channels=[cv2.UMat.get(ch) for ch in channels]\n",
    "    print(np.array(channels).shape)\n",
    "    #channels=np.moveaxis(channels, 0, -1)\n",
    "    return channels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pil_bgr_to_cmyk(bgr):\n",
    "    img = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)\n",
    "    img = Image.fromarray(img)\n",
    "    img = ImageCms.profileToProfile(img, 'color_profiles/AdobeRGB1998.icc', 'color_profiles/USWebUncoated.icc', outputMode='CMYK')\n",
    "#     img = img.convert('CMYK')\n",
    "    return np.asarray(img)\n",
    "\n",
    "def to_color_channels(bgr):\n",
    "    channels=[]\n",
    "    \n",
    "    bgr=equalize_hist(bgr)\n",
    "    print(np.array(bgr).shape)\n",
    "    \n",
    "    hsv = cv2.cvtColor(bgr, cv2.COLOR_BGR2HSV)  \n",
    "    print(\"hsv\",hsv.shape)\n",
    "    channels.append(hsv[:,:,0])\n",
    "    channels.append(hsv[:,:,1])\n",
    "    print(len(channels))\n",
    "    \n",
    "    cmyk = pil_bgr_to_cmyk(bgr)\n",
    "    print(\"cmyk\",cmyk.shape)\n",
    "    channels.append(cmyk[:,:,0])\n",
    "    channels.append(cmyk[:,:,1])\n",
    "    channels.append(cmyk[:,:,2])\n",
    "    print(len(channels))\n",
    "    \n",
    "    lab = cv2.cvtColor(bgr, cv2.COLOR_BGR2LAB)\n",
    "    print(\"lab\",lab.shape)\n",
    "#     channels.append(lab[:,:,1])\n",
    "    channels.append(lab[:,:,2])\n",
    "    print(len(channels))\n",
    "    channels=[cv2.equalizeHist(ch) for ch in channels]\n",
    "    \n",
    "    return channels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "KERNEL_W_RATIO=7.0/640.0\n",
    "\n",
    "def binarize(channels):\n",
    "    channels=np.array(channels)\n",
    "    img=np.moveaxis(channels, 0, -1)\n",
    "    \n",
    "    os.chdir(PROJECT_PATH)\n",
    "    color_clfs=[load('naive_bayes_clf_{}.joblib'.format(color)) for color in COLORS]\n",
    "    color_clf_map=dict(zip(COLORS, color_clfs))\n",
    "    \n",
    "    s0, s1, s2 = img.shape\n",
    "    flat_img = np.reshape(img, (s0*s1, s2))\n",
    "    flat_binaries=[clf.predict(flat_img) for clf in color_clfs]\n",
    "    img_binaries=[np.reshape(flat, (s0, s1)) for flat in flat_binaries]\n",
    "    return img_binaries\n",
    "\n",
    "def open_bin(img):\n",
    "    h,w=img.shape\n",
    "    k=int(w*KERNEL_W_RATIO)\n",
    "    kernel=(k,k)\n",
    "    kernel = np.ones(kernel, np.uint8) \n",
    "    img_erosion = cv2.erode(img, kernel, iterations=1)\n",
    "    img_dilation = cv2.dilate(img_erosion, kernel, iterations=1)\n",
    "    return img_dilation\n",
    "\n",
    "def close_bin(img):\n",
    "    h,w=img.shape\n",
    "    k=int(w*KERNEL_W_RATIO)\n",
    "    kernel=(k,k)\n",
    "    print(kernel)\n",
    "    kernel = np.ones(kernel, np.uint8) \n",
    "    img_dilation = cv2.dilate(img, kernel, iterations=1)\n",
    "    img_erosion = cv2.erode(img_dilation, kernel, iterations=1)\n",
    "    return img_erosion\n",
    "\n",
    "def morpho(img):\n",
    "    h,w=img.shape\n",
    "    k=int(w*KERNEL_W_RATIO)\n",
    "    kernel=(k,k)\n",
    "    return np.bool_(close_bin(open_bin(np.uint8(img), kernel), kernel))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(240, 426, 3)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "img = imread_resized(\"wooden_folia_6\", \"0001.jpg\", None)\n",
    "print(img.shape)\n",
    "cv2.imwrite(\"visualized/img.jpg\", img)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(10, 240, 426)\n"
     ]
    }
   ],
   "source": [
    "channels=convert_to_gradient_channels(img)\n",
    "\n",
    "for i in range(len(channels)):\n",
    "    cv2.imwrite(\"visualized/channel_{}.jpg\".format(i), channels[i])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(160, 213, 3)\n",
      "hsv (160, 213, 3)\n",
      "2\n",
      "cmyk (160, 213, 4)\n",
      "5\n",
      "lab (160, 213, 3)\n",
      "6\n"
     ]
    }
   ],
   "source": [
    "channels=to_color_channels(img)\n",
    "\n",
    "# for i in range(len(channels)):\n",
    "#     cv2.imwrite(\"visualized/color_channel_{}.jpg\".format(i), channels[i])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(160, 213)\n",
      "[[  0   0   0 ...   0   0   0]\n",
      " [  0   0   0 ...   0   0   0]\n",
      " [  0   0   0 ...   0   0   0]\n",
      " ...\n",
      " [  0   0   0 ...   0   0 255]\n",
      " [  0   0   0 ... 255   0   0]\n",
      " [  0   0 255 ...   0   0   0]]\n",
      "(2, 2)\n",
      "[[0 0 0 ... 0 0 0]\n",
      " [0 0 0 ... 0 0 0]\n",
      " [0 0 0 ... 0 0 0]\n",
      " ...\n",
      " [0 0 0 ... 0 0 0]\n",
      " [0 0 0 ... 0 0 0]\n",
      " [0 0 0 ... 0 0 0]]\n",
      "(2, 2)\n",
      "[[  0   0   0 ... 255   0 255]\n",
      " [255   0   0 ... 255   0   0]\n",
      " [  0   0   0 ... 255   0 255]\n",
      " ...\n",
      " [  0   0   0 ... 255   0   0]\n",
      " [  0   0 255 ...   0   0   0]\n",
      " [  0   0   0 ... 255 255 255]]\n",
      "(2, 2)\n"
     ]
    }
   ],
   "source": [
    "binary_imgs = binarize(channels)\n",
    "print(binary_imgs[0].shape)\n",
    "for i in range(len(binary_imgs)):\n",
    "    img=np.uint8(binary_imgs[i])*255\n",
    "    print(img)\n",
    "    cv2.imwrite(\"visualized/binary_img_{}.jpg\".format(i), img)\n",
    "    closed=close_bin(img)\n",
    "    cv2.imwrite(\"visualized/closed_img_{}.jpg\".format(i), closed)\n",
    "    opened=open_bin(closed)\n",
    "    cv2.imwrite(\"visualized/opened_img_{}.jpg\".format(i), opened)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "k=3\n",
    "h=18*k\n",
    "w=100*k\n",
    "window=img[0:h, 0:w]\n",
    "cv2.imwrite(\"visualized/window.jpg\", window)\n",
    "cv2.rectangle(img, (0,0), (w,h), (0, 255, 0), 4)\n",
    "cv2.imwrite(\"visualized/img_window.jpg\", img)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
